import sqlite3

import util.lib as lib

import pandas as pd
import numpy as np


# 连接到SQLite数据库（如果不存在则创建）
conn = sqlite3.connect('data/wugu.db')

# 创建表
conn.execute('''CREATE TABLE IF NOT EXISTS wugu (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    trade_code TEXT NOT NULL,
                    name  TEXT NOT NULL)''')


def update_data():
    # 打乱顺序
    df = lib.lib_get_all_stock().sample(frac=1)
    df = df[df['category'] == '股票']
    
    # 插入记录
    conn.executemany("INSERT INTO wugu (trade_code, name) VALUES (?, ?)", df[['trade_code', 'name']].to_records(index=False))
    conn.commit()

    print("更新数据")


def get_data_from_db(n = 100):
    # 从数据库中读取数据并转换为DataFrame
    return pd.read_sql_query("SELECT * FROM wugu limit " + str(n) , conn)

def delete_data(df):
    conn.execute("Delete from wugu where id in (" + ','.join(str(x) for x in df.id.values) + ")")
    conn.commit()



def get_data(n = 100):
    # 从数据库中读取数据并转换为DataFrame
    df = get_data_from_db(n)
    if len(df) == 0:
        update_data()
        df = get_data_from_db(n)

    print(df)
    delete_data(df)
    return df